{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "b8a6d60a-2951-44dc-a3cc-5967140b8ba9",
   "metadata": {},
   "source": [
    "读取与存储"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e75d9516-7282-4d81-b8a3-85bba2670fc8",
   "metadata": {},
   "source": [
    "我们有时需要把模型部署到不同的设备中，我们需要把内存中训练好的模型参数存储在硬盘中供后续读取使用"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9ff402dd-011c-4f31-afea-79f27521920e",
   "metadata": {},
   "source": [
    "1.1读写TENSOR"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0ef9f642-aca8-4bb5-b243-1ccf84807393",
   "metadata": {},
   "source": [
    "我们可以直接使用save与load函数分别存储和读取Tensor。使⽤ save 可以保存各种对象,包括模型、张\n",
    "量和字典等。⽽ laod 使⽤pickle unpickle⼯具将pickle的对象⽂件反序列化为内存。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "e609f79f-4da9-4449-9dfb-93937bc12e6f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch import nn\n",
    "\n",
    "#保存模型\n",
    "x = torch.ones(3)\n",
    "torch.save(x, 'x.pt')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "d426a8d3-4896-49ce-a132-f36c34e67c8c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([1., 1., 1.])"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#加载模型\n",
    "x2 = torch.load('x.pt', weights_only=True)\n",
    "x2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "ebb71f52-7abe-448a-8ac7-d92361aec25a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[tensor([1., 1., 1.]), tensor([0., 0., 0., 0.])]"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#存储一个Tensor列表并读回内存\n",
    "y = torch.zeros(4)\n",
    "torch.save([x, y], 'xy.pt')\n",
    "xy_list = torch.load('xy.pt',weights_only=True)\n",
    "xy_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "7195c045-9515-4a93-b62f-b90bf19b391b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'x': tensor([1., 1., 1.]), 'y': tensor([0., 0., 0., 0.])}"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#存储并读取一个从字符串映射到Tensor的字典\n",
    "torch.save({'x': x, 'y': y}, 'xy_dict.pt')\n",
    "xy = torch.load('xy_dict.pt',weights_only=True)\n",
    "xy"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e61eacbe-a4d9-4d77-9bac-53bf07baab0a",
   "metadata": {},
   "source": [
    "1.2读写模型"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e0158fe5-ee37-4fc0-84a0-533dfba2436c",
   "metadata": {},
   "source": [
    "Module的可学习参数（权重与偏差），模块模型包含在参数中(通过model.parameters() 访问)\n",
    "state_dict 是⼀个从参数名称隐射到参数 Tesnor 的字典对象。"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "01b9342a-64fd-46e9-a07b-81fb7c0b40b1",
   "metadata": {},
   "source": [
    "state_dict存储了神经网络模型的所有可学习参数。保存它可以在训练结束后保存模型的训练状态以便后续的回复于推理。可以使用一个已经训练好的模型来初始化一个新的模型，不必重新训练可以直接开始推理"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "af80f5a9-c2ce-484c-819d-a164b2c08301",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "OrderedDict([('hidden.weight',\n",
       "              tensor([[-0.2357, -0.2530, -0.2187],\n",
       "                      [-0.4763,  0.1067, -0.5626]])),\n",
       "             ('hidden.bias', tensor([ 0.4828, -0.1568])),\n",
       "             ('output.weight', tensor([[ 0.5771, -0.4153]])),\n",
       "             ('output.bias', tensor([-0.1479]))])"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "class MLP(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(MLP, self).__init__()\n",
    "        self.hidden = nn.Linear(3, 2)\n",
    "        self.act = nn.ReLU()\n",
    "        self.output = nn.Linear(2, 1)\n",
    "    def forward(self, x):\n",
    "        a = self.act(self.hidden(x))\n",
    "        return self.output(a)\n",
    "net = MLP()\n",
    "#参数的字典对象\n",
    "net.state_dict()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0bcd2eb6-9da2-4a97-bf48-0fb5f1be27e6",
   "metadata": {},
   "source": [
    "只有可学习参数的层（卷积、线性层等）才有state_dict中的条目，优化器也有state_dict，其中包含关于优化器状态以及所使用的超参数信息"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "c3ff2937-54f4-4f6d-912c-0cd799fcb13e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'state': {},\n",
       " 'param_groups': [{'lr': 0.001,\n",
       "   'momentum': 0.9,\n",
       "   'dampening': 0,\n",
       "   'weight_decay': 0,\n",
       "   'nesterov': False,\n",
       "   'maximize': False,\n",
       "   'foreach': None,\n",
       "   'differentiable': False,\n",
       "   'fused': None,\n",
       "   'params': [0, 1, 2, 3]}]}"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "optimizer = torch.optim.SGD(net.parameters(), lr=0.001, momentum=0.9)\n",
    "optimizer.state_dict()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "43d52976-8c98-465e-a580-014e44e1bb50",
   "metadata": {},
   "source": [
    "保存和加载state_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8b905653-03e0-4c4d-9b44-b86aa806c5e7",
   "metadata": {},
   "outputs": [],
   "source": [
    "#保存\n",
    "torch.save(model.state_dict(), PATH) # 推荐的⽂件后缀名是pt或pth"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8da5f770-9b9f-427f-a53a-90b118d89ea1",
   "metadata": {},
   "outputs": [],
   "source": [
    "#加载\n",
    "model = TheModelClass(*args, **kwargs)\n",
    "model.load_state_dict(torch.load(PATH))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "ffe0f53a-eac0-4e8d-9da2-3f9d5946742f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[0.2215],\n",
      "        [0.1759]], grad_fn=<AddmmBackward0>)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "tensor([[True],\n",
       "        [True]])"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "X = torch.randn(2, 3)\n",
    "Y = net(X)\n",
    "PATH = \"./net.pt\"\n",
    "torch.save(net.state_dict(), PATH)\n",
    "net2 = MLP()\n",
    "net2.load_state_dict(torch.load(PATH, weights_only=True))\n",
    "Y2 = net2(X)\n",
    "print(Y2)\n",
    "Y2 == Y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9cc50712-79de-433f-a060-1f7ebcbbe2db",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:pytorch]",
   "language": "python",
   "name": "deep_learning"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
